import os
import requests
import time
import random
import json
from urllib.parse import urljoin
from datetime import datetime, timedelta
import sys
import logging

"""
   深交所定期财报
   用法1: python szse.py (使用默认3年范围)
   用法2: python szse.py 开始日期 结束日期 (格式: YYYY-MM-DD)
"""

def get_next_log_filename(base_filename):
    """获取下一个可用的日志文件名，添加_n后缀"""
    counter = 1
    filename = f"{base_filename}.log"

    while os.path.exists(os.path.join("logs", "szse", filename)):
        filename = f"{base_filename}_{counter}.log"
        counter += 1

    return filename

def setup_logging():
    """配置日志记录，同时输出到控制台和文件"""
    # 创建日志目录（如果需要）
    log_dir = os.path.join(os.getcwd(), "logs","szse")
    os.makedirs(log_dir, exist_ok=True)

    # 基础文件名（不含后缀）
    base_filename = f"szse{datetime.now().strftime('%Y%m%d')}"

    # 获取下一个可用的日志文件名
    log_filename = get_next_log_filename(base_filename)
    log_filepath = os.path.join(log_dir, log_filename)

    # 配置日志
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(log_filepath, encoding='utf-8'),
            logging.StreamHandler(sys.stdout)
        ]
    )
    return logging.getLogger()

# 初始化日志
logger = setup_logging()

# 请求配置
url = "https://www.szse.cn/api/disc/announcement/annList"

# 随机User-Agent生成器
def get_random_user_agent():
    user_agents = [
        "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36",
        "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:89.0) Gecko/20100101 Firefox/89.0",
        "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36",
        "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/14.1.1 Safari/605.1.15"
    ]
    return random.choice(user_agents)

# 伪装UA
headers = {
    "User-Agent": get_random_user_agent(),
    "Accept": "application/json, text/javascript, */*; q=0.01",
    "Accept-Language": "zh-CN,zh;q=0.9",
    "Accept-Encoding": "gzip, deflate, br",
    "Content-Type": "application/json",
    "Origin": "https://www.szse.cn",
    "Referer": "https://www.szse.cn/disclosure/announcement/company/index.html",
    "X-Requested-With": "XMLHttpRequest",
    "Connection": "keep-alive"
}

def get_date_range():
    """获取日期范围"""
    # 如果没有提供参数，使用默认的3年范围
    if len(sys.argv) == 1:
        today = datetime.now()
        three_years_ago = today - timedelta(days=3*365)
        first_date = three_years_ago.replace(month=1, day=1).strftime("%Y-%m-%d")
        second_date = today.strftime("%Y-%m-%d")
        return first_date, second_date

    # 如果提供了参数但数量不对
    if len(sys.argv) != 3:
        logger.error("错误：参数数量不正确")
        logger.info("用法1: python szse.py (使用默认3年范围)")
        logger.info("用法2: python szse.py 开始日期 结束日期 (格式: YYYY-MM-DD)")
        sys.exit(1)

    # 验证日期格式
    try:
        start_date = datetime.strptime(sys.argv[1], "%Y-%m-%d").strftime("%Y-%m-%d")
        end_date = datetime.strptime(sys.argv[2], "%Y-%m-%d").strftime("%Y-%m-%d")

        # 检查开始日期是否早于结束日期
        if datetime.strptime(start_date, "%Y-%m-%d") > datetime.strptime(end_date, "%Y-%m-%d"):
            logger.error("错误：开始日期不能晚于结束日期")
            sys.exit(1)

        return start_date, end_date

    except ValueError as e:
        logger.error(f"错误：日期格式不正确 - {e}")
        logger.info("请使用 YYYY-MM-DD 格式的日期，例如: 2025-01-01")
        sys.exit(1)

payload = {
    # "bigCategoryId": ["010301"],
    "seDate": get_date_range(),
    "channelCode": ["fixed_disc"],
    "pageSize": 50,
    "pageNum": 1,
    # "stock": ["000034", "000048", "000069", "000333", "000408", "000528", "000539", "000550", "000555", "000629", "000663", "000738", "000756", "000776", "000792", "000800", "000898", "000901", "000919", "000922", "000977", "001277", "001289", "001308", "002024", "002030", "002096", "002106", "002124", "002140", "002202", "002236", "002362", "002401", "002460", "002467", "002534", "002552", "002672", "002703", "002723", "002732", "003035", "300007", "300019", "300101", "300183", "300188", "300301", "300363", "300373", "300406", "300426", "300446", "300455", "300457", "300476", "300479", "300522", "300575", "300620", "300657", "300674", "300699", "300748", "300782", "300910", "301004", "301008", "301011", "301016", "301059", "301113", "301121"]
}

# 基础下载URL
base_download_url = "https://disc.static.szse.cn/download"

def random_delay():
    delay_pattern = random.choice([
        random.uniform(0.5, 1),
        random.uniform(2, 3),
        random.uniform(4, 6)
    ])
    time.sleep(delay_pattern)

    if random.randint(1, 10) == 1:
        time.sleep(random.uniform(7, 8))

def download_file(url, filepath):
    """下载文件并保存到指定路径"""
    try:
        random_delay()
        response = requests.get(url, stream=True, headers=headers)
        response.raise_for_status()

        os.makedirs(os.path.dirname(filepath), exist_ok=True)

        # 临时文件路径
        temp_filepath = filepath + '.tmp'

        try:
            with open(temp_filepath, 'wb') as f:
                for chunk in response.iter_content(chunk_size=8192):
                    if chunk:
                        f.write(chunk)

            # 下载完成后重命名临时文件为正式文件
            os.rename(temp_filepath, filepath)
            logger.info(f"成功下载: {filepath}")

        except Exception as e:
            # 如果下载过程中出错，删除临时文件
            if os.path.exists(temp_filepath):
                os.remove(temp_filepath)
            raise e

    except Exception as e:
        # 如果文件已存在但下载失败，删除可能不完整的文件
        if os.path.exists(filepath):
            try:
                os.remove(filepath)
                logger.warning(f"已删除不完整文件: {filepath}")
            except Exception as delete_error:
                logger.error(f"删除文件失败 {filepath}: {delete_error}")

        logger.error(f"下载失败 {url}: {e}")

def process_announcements():
    """处理公告数据"""
    try:
        # 第一次请求获取总数量
        response = requests.post(url, headers=headers, data=json.dumps(payload))
        response.raise_for_status()
        data = response.json()

        total_count = data.get("announceCount", 0)
        logger.info(f"总公告数: {total_count}")

        if total_count == 0:
            logger.info("没有找到公告")
            return

        # 计算总页数
        page_size = payload["pageSize"]
        total_pages = (total_count + page_size - 1) // page_size

        # 处理每一页
        for page_num in range(1, total_pages + 1):
            logger.info(f"\n正在处理第 {page_num}/{total_pages} 页...")
            payload["pageNum"] = page_num

            random_delay()
            response = requests.post(url, headers=headers, data=json.dumps(payload))
            response.raise_for_status()
            page_data = response.json()

            # 处理每条公告
            for announcement in page_data.get("data", []):
                try:
                    sec_code = announcement.get("secCode", [""])[0]
                    sec_name = announcement.get("secName", [""])[0]
                    title = announcement.get("title", "")
                    publish_time = announcement.get("publishTime", "")
                    publish_date = datetime.strptime(publish_time, "%Y-%m-%d %H:%M:%S")
                    formatted_date = publish_date.strftime("%Y-%m-%d")
                    attach_path = announcement.get("attachPath", "")

                    if not all([sec_code, sec_name, title, formatted_date, attach_path]):
                        logger.warning(f"跳过不完整数据: {announcement}")
                        continue

                    # 创建目录
                    dir_name = f"{sec_code}_{sec_name}"
                    dir_name = "".join(c for c in dir_name if c not in '\\/:*?"<>|')  # 移除非法字符
                    base_dir = os.path.join(os.getcwd(), "report", "szse")  # 基础目录结构
                    dir_path = os.path.join(base_dir, dir_name)  # 完整目录路径

                    # 创建文件名
                    file_name = f"{title}_{formatted_date}.pdf"
                    file_name = "".join(c for c in file_name if c not in '\\/:*?"<>|')  # 移除非法字符
                    file_path = os.path.join(dir_path, file_name)

                    # 检查文件是否已存在
                    if os.path.exists(file_path):
                        logger.info(f"文件已存在，跳过下载: {file_path}")
                        continue  # 跳过当前公告，处理下一条

                    # 下载文件
                    download_url = urljoin(base_download_url, attach_path)
                    download_file(download_url, file_path)

                except Exception as e:
                    logger.error(f"处理公告时出错: {e}")
                    continue

    except Exception as e:
        logger.error(f"处理过程中出错: {e}")
        raise  # 重新抛出异常以便主程序捕获

if __name__ == "__main__":
    try:
        logger.info("="*50)
        logger.info(f"程序开始运行 - {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        logger.info(f"命令行参数: {sys.argv}")

        process_announcements()

        logger.info("处理完成")
        logger.info(f"程序结束运行 - {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        logger.info("="*50)
    except KeyboardInterrupt:
        logger.warning("\n用户中断操作")
        sys.exit(1)
    except Exception as e:
        logger.error(f"程序运行出错: {e}")
        sys.exit(1)